# -*- coding:utf-8 -*-

"""
------------------------------------
# @FileName    :main.py
# @Time        :2025/10/14 8:48
# @Author      :hanyj
# @description : npu利用率查询（AICORE）：npu-smi info watch
------------------------------------
"""
import asyncio
import aiohttp
import time
import os
import yaml
from typing import List, Dict
import re

# ================== 配置区 ==================
INTERVAL = 30  # 定时间隔（秒）

# 加载配置文件
CONFIG = {}
try:
    current_dir = os.path.dirname(__file__)
    config_path = os.path.join(current_dir, "config.yaml")
    if os.path.exists(config_path):
        with open(config_path, "r", encoding="utf-8") as f:
            CONFIG = yaml.safe_load(f) or {}
            if not isinstance(CONFIG, dict):
                CONFIG = {}
except Exception as e:
    print(f"加载 config.yaml 失败：{e}")

# 从配置文件读取超时时间，默认300秒（5分钟）
TIMEOUT = CONFIG.get("LLM_TIMEOUT", 300)

# 从配置文件获取基础URL
LLM_QWEN3_14B_API_URL = CONFIG.get("LLM_QWEN3_14B_API_URL", "http://10.18.214.83:1040/v1")

# 每个 API 的配置：包含 URL、并发次数、可选 body/headers 等
API_CONFIGS: List[Dict] = [
    {
        "name": "qwen3-14b",
        "url": f"{LLM_QWEN3_14B_API_URL}/chat/completions",
        "concurrent": 1,
        "json_data": {
            "model": "qwen3-14b",
            "messages": [{"role": "system", "content": "/no_think"}, {"role": "user", "content": "请给我讲一个笑话！"}],
            "max_tokens": 1024,
            "temperature": 0.0,
            "stream": False
        }
    }
]

# 全局 headers（如 token）
HEADERS = {
    "Content-Type": "application/json",
}


# ==========================================

async def call_single_api(session: aiohttp.ClientSession, config: Dict, system_prompt=None, user_prompt=None):
    """调用单个 API 一次（带重试机制）"""
    name = config["name"]
    url = config["url"]
    json_data = config.get("json_data", None)
    headers = {**HEADERS, **config.get("headers", {})}

    # 覆盖 messages 中的 content（若提供）
    if json_data and isinstance(json_data, dict):
        messages = json_data.get("messages")
        if isinstance(messages, list) and len(messages) >= 2:
            if system_prompt is not None and isinstance(messages[0], dict):
                messages[0]["content"] = system_prompt
            if user_prompt is not None and isinstance(messages[1], dict):
                messages[1]["content"] = user_prompt

    # 重试机制：最多重试3次
    for attempt in range(3):
        try:
            async with session.post(
                    url=url,
                    json=json_data,
                    headers=headers,
                    timeout=aiohttp.ClientTimeout(total=TIMEOUT)
            ) as response:
                status = response.status
                try:
                    data = await response.json()
                except Exception:
                    text = await response.text()
                    print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] ❌ {name} 非JSON({status}): {text[:500]}")
                    return
                if status >= 400:
                    print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] ❌ {name} HTTP{status}: {str(data)[:1000]}")
                    return None
                if isinstance(data, dict) and data.get("error"):
                    print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] ❌ {name} 业务错误: {data['error']}")
                    return None
                msg = None
                if isinstance(data, dict):
                    choices = data.get("choices") or []
                    if choices and isinstance(choices[0], dict):
                        msg = (choices[0].get("message") or {}).get("content")
                    if msg is None:
                        msg = data.get("content")
                result_text = str(msg if msg is not None else data)
                # 清理<think>标签后再打印，避免日志混乱
                clean_text = re.sub(r'<think>.*?</think>', '', result_text, flags=re.DOTALL).strip()
                print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] ✅ {name} 回复: {clean_text[:2000]}")
                return result_text
        except asyncio.TimeoutError:
            if attempt < 2:  # 还有重试机会
                print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] ⚠️ {name} 超时，第{attempt+1}次重试...")
                await asyncio.sleep(2 ** attempt)  # 指数退避：2秒、4秒
                continue
            else:
                print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] ❌ {name} 超时（已重试3次）")
        except Exception as e:
            if attempt < 2:  # 还有重试机会
                print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] ⚠️ {name} 错误，第{attempt+1}次重试: {e}")
                await asyncio.sleep(2 ** attempt)
                continue
            else:
                print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] ❌ {name} 错误（已重试3次）: {e}")
    return None


async def call_api_with_concurrency(session: aiohttp.ClientSession, config: Dict, system_prompt=None, user_prompt=None, concurrent_override=None):
    """对单个 API 发起指定次数的并发调用"""
    name = config["name"]
    concurrent = concurrent_override if concurrent_override is not None else config.get("concurrent", 1)
    print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] ➡️  开始并发调用 {name} x{concurrent} 次")

    tasks = [call_single_api(session, config, system_prompt, user_prompt) for _ in range(concurrent)]
    results = await asyncio.gather(*tasks)
    return results


# ===== 最小改动：增加入口以触发并发调用 =====
async def _run_all_configs_once(name=None, system_prompt=None, user_prompt=None, concurrent=None):
    # 按名称筛选（为 None 时运行全部）
    if name is None:
        name = "qwen3-14b"  # 默认使用qwen3-14b
    selected_configs = [cfg for cfg in API_CONFIGS if cfg.get("name") == name]
    if not selected_configs:
        print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] ⚠️ 未找到名称为 {name} 的配置")
        return

    # async with aiohttp.ClientSession() as session: 这是不限制
    # 限制同一主机的并发连接，避免击穿模型服务
    connector = aiohttp.TCPConnector(limit_per_host=1)
    async with aiohttp.ClientSession(connector=connector) as session:
        all_results = await asyncio.gather(*(call_api_with_concurrency(session, cfg, system_prompt, user_prompt, concurrent_override=concurrent) for cfg in selected_configs))
        # 扁平化：List[List[Optional[str]]] -> List[str]
        flat = []
        for sub in all_results:
            for item in (sub or []):
                if item is not None:
                    flat.append(item)
        return flat


if __name__ == "__main__":
    asyncio.run(_run_all_configs_once())

